import datetime
import gradio as gr
import pymongo as pm
from util_mongo import enqueue, get_sorted_by_key
from common import MONGODB_NAME, VALUE, KEY, IO_PREFIX
from util import uuid, merge_dialog_in_and_out, translate_ns
import time
from transformers import AutoTokenizer, AutoModel
from http.cookies import SimpleCookie

# 连接Mongodb
mongo = pm.MongoClient('127.0.0.1', 27017, serverSelectionTimeoutMS=3000)
mdb = mongo[MONGODB_NAME]

print('-------------------------------------------------------')
print('正在加载模型……')
# model_name = "THUDM/chatglm2-6b-int4"
model_name = "/root/.cache/huggingface/hub/models--THUDM--chatglm2-6b-int4/snapshots/66ecaf1db3a5085714e133357ea4824b69698743"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(model_name,trust_remote_code=True).cuda()
model = model.eval()
print('模型已经加载完毕。')


def chat_once(xinput, username):
    """输入文本，用模型返回输出。"""
    # 获取聊天历史
    xlog = get_history(username)
    
    # 得到输出
    print('-------------history-----------------')
    for i, (xin, xout) in enumerate(xlog):
        print(i, '>>>>', xin)
        print(i, '<<<<', xout)
    print('-------------this turn---------------')
    print('>>>>', '>>>>', xinput)
    xoutput, history = model.chat(tokenizer, xinput, history=xlog)
    print('<<<<', '<<<<', xoutput)
    
    # 输入存入mongodb
    ts = time.time_ns()
    enqueue(mdb, 'dialog_in', username, ts, {
        VALUE: xinput,
        'io': 'i',
    })
    
    # 输出存入mongodb
    ts = time.time_ns()
    enqueue(mdb, 'dialog_out', username, ts, {
        VALUE: xoutput,
        'io': 'o',
    })

    # 获取聊天历史
    xlog = get_history(username, more_info=True)

    # 返回空字符串清空输入文本框，返回聊天历史
    return '', xlog


def get_history(username, more_info=False):
    """
    获取聊天历史

    把历史组合成对子的列表，例如：
    in: xxx
    out: yyy
    in: zzzz
    in: aaaa
    out: bbbb
    out: kkkk
    =>
    [
        [xxx, yyy],
        [zzzz, None],
        [aaaa, bbbb],
        [None, kkkk],
    ]

    """
    # 获取输入历史
    rows_in = get_sorted_by_key(mdb, 'dialog_in', username, limit=6, is_keep_others=False)
    # 获取输出历史
    rows_out = get_sorted_by_key(mdb, 'dialog_out', username, limit=6, is_keep_others=False)
    # 合并输入历史和输出历史
    rows = merge_dialog_in_and_out(rows_in, rows_out)  # DESC
    # 按时间正序排列
    rows = rows[::-1]  # ASC

    if not len(rows):
        # 没有历史则返回空列表
        log = []
    else:
        # 把聊天历史组合成对子的列表
        log = []
        pair = [None, None]
        for row in rows:
            text = row[VALUE]
            
            if not more_info:
                xstr = text
            else:
                timestamp = row[KEY]
                dt_str = translate_ns(timestamp)
                xstr = IO_PREFIX.get(row["io"], '?') + ': '
                xstr += f'[ {dt_str} ] '

                xstr += text

            if 'i' == row['io']:
                # 如果是输入

                # 如果对子里面已经有输入或输出了，则新建对子
                if pair[0] is not None or pair[1] is not None:
                    log.append(pair)
                    pair = [xstr, None]
                    continue

                # 输入放入对子的第一个元素
                pair[0] = xstr

            elif 'o' == row['io']:
                # 如果是输出

                # 如果对子里面已经有输出了，则新建对子
                if pair[1] is not None:
                    log.append(pair)
                    pair = [None, xstr]
                    continue

                # 输出放入对子的第2个元素
                pair[1] = xstr
        log.append(pair)

    return log


def embed_user_id_in_state_comp(request: gr.Request):
    """
    页面加载的事件handler

    从HTTP request中获取cookie来确定用户名，如果没有则生成一个
    :param request: HTTP request data
    :return: 将用户名发送给浏览器中存放用户名的组件
    """
    xuuid = uuid()

    #从cookie中获取username
    username = None
    try:
        cookie_str = request.headers.cookie
        print('cookie_str:', cookie_str)
        cookie = SimpleCookie()
        cookie.load(cookie_str)
        username = cookie.get('username', None)
        print('user:', username)
        username = username.value
    except Exception:
        try:
            username = request.cookies['username']
        except Exception:
            pass

    # 获取不到，新建一个username
    if username is None:
        print('new user!')
        username = f'u{xuuid}'

    # print username
    print('user:', username)

    xlog = get_history(username, more_info=True)

    return username, username, username, xlog


# 用户名写入cookie的JS
js_username2cookie = """
    async () => {{
            setTimeout(() => {{
                var xnew = document.querySelector('#username.prose').innerHTML;
                console.log('username', xnew)
                if ('' == xnew) {
                    return;
                }
                document.cookie = 'username=' + xnew + '; max-age=' + 3600*24*30*6;
            }}, 1);
        }}
"""

# 构建Gradio界面
with gr.Blocks(analytics_enabled=False) as demo:

    # 用户名相关控件
    cmp_username_state = gr.State()
    cmp_username_html = gr.HTML(visible=False, elem_id='username')
    cmp_username_display = gr.Textbox(interactive=False, label='用户名')

    # 输入文本框
    cmp_input = gr.Textbox(interactive=True, label='输入')

    # 输入按钮
    cmp_input_button = gr.Button('发送', )

    # 聊天历史
    cmp_chatbot = gr.Chatbot().style(height=350)

    # 输入按钮点击事件
    cmp_input_button.click(chat_once, [cmp_input, cmp_username_state], [
        cmp_input,
        cmp_chatbot,
    ], queue=False)

    # 加载时获取或设置cookie，并展示聊天历史
    demo.load(embed_user_id_in_state_comp, None, [
        cmp_username_state,
        cmp_username_display,
        cmp_username_html,
        cmp_chatbot,
    ], queue=False)
    # cmp_username_html发生变化时，触发JS，把生成的username放入页面cookie
    cmp_username_html.change(None, None, None, _js=js_username2cookie, queue=False)

# 使用queue（目前这个例子太简单，没有用到queue，后面会用到。）
demo.queue(
    concurrency_count=10,
    status_update_rate='auto',
    # status_update_rate=0.02,
)

# 启动Gradio
# demo.launch(server_name='0.0.0.0', server_port=6006, share=True, debug=True)  # 带内网穿透
demo.launch(server_name='0.0.0.0', server_port=6006, debug=True)
